#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/6/5 22:44
# @USER    : Shengji He
# @File    : MergeKSortedLists.py
# @Software: PyCharm
# @Version  : Python-
# @TASK:
import heapq
from typing import List


# Definition for singly-linked list.
class ListNode:
    def __init__(self, x):
        self.val = x
        self.next = None


class Solution:
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        """
        Merge k sorted linked lists and return it as one sorted list. Analyze and describe its complexity.

        Example:

        Input:

        [
            1->4->5,

            1->3->4,

            2->6

        ]

        Output: 1->1->2->3->4->4->5->6

        :param lists: List[ListNode]
        :return: ListNode
        """
        # def mergeTwo(l1: ListNode, l2: ListNode):
        #     output = ListNode(0)
        #     node = output
        #     while l1 or l2:
        #         if l1 is None and l2 is not None:
        #             node.next = l2
        #             break
        #         elif l1 is not None and l2 is None:
        #             node.next = l1
        #             break
        #         else:
        #             if l1.val > l2.val:
        #                 node.next = ListNode(l2.val)
        #                 l2 = l2.next
        #             else:
        #                 node.next = ListNode(l1.val)
        #                 l1 = l1.next
        #             node = node.next
        #     return output.next
        # return reduce(mergeTwo, lists) if len(lists) > 1 else lists
        if len(lists) < 2:
            return [[], *lists][len(lists)]
        else:
            min_heaq = []
            ans = node = ListNode(0)
            for i in range(len(lists)):
                if lists[i]:
                    heapq.heappush(min_heaq, (lists[i].val, i))
            while min_heaq:
                val, idx = heapq.heappop(min_heaq)
                node.next = ListNode(val)
                node = node.next
                lists[idx] = lists[idx].next
                if lists[idx]:
                    heapq.heappush(min_heaq, (lists[idx].val, idx))
            return ans.next


if __name__ == '__main__':
    S = Solution()
    input1 = [1, 4, 5]
    input2 = [1, 3, 4]
    input3 = [2, 6]
    print(input1)
    print(input2)
    print(input3)
    l1 = ListNode(input1.pop(0))
    l2 = ListNode(input2.pop(0))
    l3 = ListNode(input3.pop(0))
    temp1 = l1
    for element in input1:
        temp1.next = ListNode(element)
        temp1 = temp1.next
    temp2 = l2
    for element in input2:
        temp2.next = ListNode(element)
        temp2 = temp2.next
    temp3 = l3
    for element in input3:
        temp3.next = ListNode(element)
        temp3 = temp3.next
    output = S.mergeKLists([l1, l2, l3])
    print('done')
